/*
 * Copyright (c) 2024 OpenHiTLS. All rights reserved.
 * Licensed under the OpenHiTLS license.
 *
 * TLS/SSL JNI Implementation for HiTLS4J
 * Minimal implementation to verify basic functionality.
 */

#include <jni.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <stdbool.h>

// OpenHiTLS TLS headers
#include "hitls.h"
#include "hitls_config.h"
#include "hitls_error.h"
#include "hitls_type.h"
#include "hitls_crypt_init.h"
#include "hitls_cert_init.h"
#include "hitls_cert.h"
#include "hitls_sni.h"
#include "bsl_err.h"
#include "bsl_uio.h"
#include "bsl_sal.h"

// JNI header
#include "org_openhitls_core_CryptoNative.h"

// TLS initialization flag
static int g_tlsInitialized = 0;

// Custom Java BIO method for bidirectional memory I/O
static BSL_UIO_Method *g_javaBioMethod = NULL;

// ==================== Custom Java BIO Implementation ====================

// Java BIO data structure with separate input/output buffers
typedef struct {
    uint8_t *inBuf;       // Buffer for incoming network data (Java writes, HITLS reads)
    size_t inBufSize;
    size_t inBufPos;      // Read position
    size_t inBufLen;      // Write position (amount of data)

    uint8_t *outBuf;      // Buffer for outgoing network data (HITLS writes, Java reads)
    size_t outBufSize;
    size_t outBufPos;     // Read position
    size_t outBufLen;     // Write position (amount of data)
} JavaBioData;

// Read callback - HITLS reads from input buffer (incoming network data)
static int32_t javaBioRead(BSL_UIO *uio, void *buf, uint32_t len, uint32_t *readLen) {
    JavaBioData *bioData = (JavaBioData *)BSL_UIO_GetCtx(uio);
    if (bioData == NULL || bioData->inBuf == NULL) {
        *readLen = 0;
        return BSL_UIO_IO_EXCEPTION;
    }

    size_t available = bioData->inBufLen - bioData->inBufPos;
    if (available == 0) {
        *readLen = 0;
        // For non-blocking I/O with no data available:
        // Return success with 0 bytes read - HITLS will check and return HITLS_WANT_READ
        return BSL_SUCCESS;
    }

    size_t toRead = len < available ? len : available;
    memcpy(buf, bioData->inBuf + bioData->inBufPos, toRead);
    bioData->inBufPos += toRead;
    *readLen = (uint32_t)toRead;

    // Compact buffer if fully consumed
    if (bioData->inBufPos >= bioData->inBufLen) {
        bioData->inBufPos = 0;
        bioData->inBufLen = 0;
    }

    return BSL_SUCCESS;
}

// Write callback - HITLS writes to output buffer (outgoing network data)
static int32_t javaBioWrite(BSL_UIO *uio, const void *buf, uint32_t len, uint32_t *writeLen) {
    JavaBioData *bioData = (JavaBioData *)BSL_UIO_GetCtx(uio);
    if (bioData == NULL) {
        //fprintf(stderr, "javaBioWrite: bioData is NULL\n");
        *writeLen = 0;
        return BSL_UIO_IO_EXCEPTION;
    }

    //fprintf(stderr, "javaBioWrite: writing %u bytes\n", len);

    // Compact buffer if needed
    if (bioData->outBufPos > 0 && bioData->outBufPos >= bioData->outBufLen) {
        bioData->outBufPos = 0;
        bioData->outBufLen = 0;
    }

    // Check if we need to expand the buffer
    size_t needed = bioData->outBufLen + len;
    if (needed > bioData->outBufSize) {
        size_t newSize = bioData->outBufSize * 2;
        if (newSize < needed) newSize = needed;
        if (newSize < 32768) newSize = 32768;
        uint8_t *newBuf = (uint8_t *)realloc(bioData->outBuf, newSize);
        if (newBuf == NULL) {
            *writeLen = 0;
            return BSL_INTERNAL_EXCEPTION;
        }
        bioData->outBuf = newBuf;
        bioData->outBufSize = newSize;
    }

    memcpy(bioData->outBuf + bioData->outBufLen, buf, len);
    bioData->outBufLen += len;
    *writeLen = len;

    //fprintf(stderr, "javaBioWrite: outBufLen now %zu\n", bioData->outBufLen);

    return BSL_SUCCESS;
}

// Ctrl callback
static int32_t javaBioCtrl(BSL_UIO *uio, int32_t cmd, int32_t larg, void *parg) {
    JavaBioData *bioData = (JavaBioData *)BSL_UIO_GetCtx(uio);
    if (bioData == NULL) {
        return BSL_SUCCESS;  // Return success for most ctrl commands
    }

    switch (cmd) {
        case BSL_UIO_FLUSH:
            return BSL_SUCCESS;
        case BSL_UIO_PENDING:
            // Pending read data (in input buffer for HITLS to read)
            return (int32_t)(bioData->inBufLen - bioData->inBufPos);
        case BSL_UIO_WPENDING:
            // Pending write data (in output buffer for Java to read)
            return (int32_t)(bioData->outBufLen - bioData->outBufPos);
        case BSL_UIO_GET_FD:
            return -1;  // No file descriptor
        default:
            return BSL_SUCCESS;
    }
}

// Create callback
static int32_t javaBioCreate(BSL_UIO *uio) {
    JavaBioData *bioData = (JavaBioData *)malloc(sizeof(JavaBioData));
    if (bioData == NULL) {
        return BSL_INTERNAL_EXCEPTION;
    }
    memset(bioData, 0, sizeof(JavaBioData));

    // Initial buffer sizes (16KB each)
    bioData->inBufSize = 16384;
    bioData->outBufSize = 16384;
    bioData->inBuf = (uint8_t *)malloc(bioData->inBufSize);
    bioData->outBuf = (uint8_t *)malloc(bioData->outBufSize);

    if (bioData->inBuf == NULL || bioData->outBuf == NULL) {
        free(bioData->inBuf);
        free(bioData->outBuf);
        free(bioData);
        return BSL_INTERNAL_EXCEPTION;
    }

    BSL_UIO_SetCtx(uio, bioData);
    BSL_UIO_SetInit(uio, true);
    return BSL_SUCCESS;
}

// Destroy callback
static int32_t javaBioDestroy(BSL_UIO *uio) {
    JavaBioData *bioData = (JavaBioData *)BSL_UIO_GetCtx(uio);
    if (bioData != NULL) {
        free(bioData->inBuf);
        free(bioData->outBuf);
        free(bioData);
        BSL_UIO_SetCtx(uio, NULL);
    }
    return BSL_SUCCESS;
}

// Get or create the Java BIO method
static const BSL_UIO_Method *getJavaBioMethod(void) {
    if (g_javaBioMethod == NULL) {
        g_javaBioMethod = BSL_UIO_NewMethod();
        if (g_javaBioMethod == NULL) {
            return NULL;
        }

        BSL_UIO_SetMethodType(g_javaBioMethod, BSL_UIO_MEM);
        BSL_UIO_SetMethod(g_javaBioMethod, BSL_UIO_CREATE_CB, (void *)javaBioCreate);
        BSL_UIO_SetMethod(g_javaBioMethod, BSL_UIO_DESTROY_CB, (void *)javaBioDestroy);
        BSL_UIO_SetMethod(g_javaBioMethod, BSL_UIO_READ_CB, (void *)javaBioRead);
        BSL_UIO_SetMethod(g_javaBioMethod, BSL_UIO_WRITE_CB, (void *)javaBioWrite);
        BSL_UIO_SetMethod(g_javaBioMethod, BSL_UIO_CTRL_CB, (void *)javaBioCtrl);
    }
    return g_javaBioMethod;
}

static int ensureTLSInitialized(void) {
    if (!g_tlsInitialized) {
        // Note: BSL/CRYPT initialization is done by crypto_native_jni.c in JNI_OnLoad
        // We only need to initialize TLS-specific components here

        // Initialize certificate methods for TLS
        HITLS_CertMethodInit();

        // Initialize crypto methods for TLS
        HITLS_CryptMethodInit();

        g_tlsInitialized = 1;
    }
    return 0;
}

// ==================== Exception Handling ====================

static void throwSSLException(JNIEnv *env, const char *message) {
    jclass cls = (*env)->FindClass(env, "javax/net/ssl/SSLException");
    if (cls != NULL) {
        (*env)->ThrowNew(env, cls, message);
        (*env)->DeleteLocalRef(env, cls);
    }
}

// ==================== SSL Context Wrapper ====================

typedef struct {
    HITLS_Config *config;
    HITLS_Ctx *ctx;
    BSL_UIO *uio;           // Memory UIO for TLS I/O
    bool isClient;
    int handshakeState;     // 0=not started, 1=in progress, 2=complete
} SSLWrapper;

// ==================== SSL Context Management ====================

JNIEXPORT jlong JNICALL Java_org_openhitls_core_CryptoNative_sslCtxNew
  (JNIEnv *env, jclass cls, jint minVersion, jint maxVersion) {

    // Ensure TLS is initialized
    int initRet = ensureTLSInitialized();
    if (initRet != 0) {
        char errorMsg[256];
        snprintf(errorMsg, sizeof(errorMsg),
                 "TLS initialization failed with error: %d", initRet);
        throwSSLException(env, errorMsg);
        return 0;
    }

    HITLS_Config *config = NULL;

    // Try generic TLS config first (supports TLS 1.2 + 1.3)
    config = HITLS_CFG_NewTLSConfig();

    if (config == NULL) {
        // Fallback to TLS 1.2 only
        config = HITLS_CFG_NewTLS12Config();
    }

    if (config == NULL) {
        char errorMsg[256];
        snprintf(errorMsg, sizeof(errorMsg),
                 "Failed to create TLS configuration (min=%d, max=%d).",
                 minVersion, maxVersion);
        throwSSLException(env, errorMsg);
        return 0;
    }

    return (jlong)config;
}

JNIEXPORT void JNICALL Java_org_openhitls_core_CryptoNative_sslCtxFree
  (JNIEnv *env, jclass cls, jlong ctxPtr) {
    if (ctxPtr != 0) {
        HITLS_Config *config = (HITLS_Config *)ctxPtr;
        HITLS_CFG_FreeConfig(config);
    }
}

JNIEXPORT jlong JNICALL Java_org_openhitls_core_CryptoNative_sslCtxSetOptions
  (JNIEnv *env, jclass cls, jlong ctxPtr, jlong options) {
    // Stub - options handling not implemented
    return options;
}

JNIEXPORT void JNICALL Java_org_openhitls_core_CryptoNative_sslCtxSetCertificate
  (JNIEnv *env, jclass cls, jlong ctxPtr, jbyteArray certData, jint format) {
    // Stub - certificate handling requires PKI integration
}

JNIEXPORT void JNICALL Java_org_openhitls_core_CryptoNative_sslCtxSetPrivateKey
  (JNIEnv *env, jclass cls, jlong ctxPtr, jbyteArray keyData, jint format) {
    // Stub - private key handling requires PKI integration
}

JNIEXPORT void JNICALL Java_org_openhitls_core_CryptoNative_sslCtxSetCipherSuites
  (JNIEnv *env, jclass cls, jlong ctxPtr, jstring cipherSuites) {
    // Stub - cipher suite configuration not implemented
}

JNIEXPORT void JNICALL Java_org_openhitls_core_CryptoNative_sslCtxSetCipherSuitesTls13
  (JNIEnv *env, jclass cls, jlong ctxPtr, jstring cipherSuites) {
    // Stub - TLS 1.3 cipher suite configuration not implemented
}

JNIEXPORT void JNICALL Java_org_openhitls_core_CryptoNative_sslCtxSetVerifyMode
  (JNIEnv *env, jclass cls, jlong ctxPtr, jint mode) {
    HITLS_Config *config = (HITLS_Config *)ctxPtr;
    if (config != NULL && mode > 0) {
        HITLS_CFG_SetClientVerifySupport(config, true);
    }
}

JNIEXPORT void JNICALL Java_org_openhitls_core_CryptoNative_sslCtxSetTrustCertificate
  (JNIEnv *env, jclass cls, jlong ctxPtr, jbyteArray certData, jint format) {
    // Stub - trust certificate handling requires PKI integration
}

// Global pointer to track the current ctx for verify callback (needed because callback doesn't receive ctx)
static __thread HITLS_Ctx *g_currentCtxForVerify = NULL;

// Trust-all verify callback - always returns 1 (success)
static int trustAllVerifyCallback(int32_t isPreverifyOk, HITLS_CERT_StoreCtx *storeCtx) {
    (void)storeCtx;
    (void)isPreverifyOk;
    // Also set verify result to success if we have the ctx
    if (g_currentCtxForVerify != NULL) {
        HITLS_SetVerifyResult(g_currentCtxForVerify, HITLS_SUCCESS);
    }
    return 1;  // Always accept
}

// Real certificate verification callback - validates certificate chain
static int realVerifyCallback(int32_t isPreverifyOk, HITLS_CERT_StoreCtx *storeCtx) {
    (void)storeCtx;
    // Return the pre-verification result from openHiTLS
    // openHiTLS has already validated the certificate chain against loaded CA certs
    if (isPreverifyOk) {
        if (g_currentCtxForVerify != NULL) {
            HITLS_SetVerifyResult(g_currentCtxForVerify, HITLS_SUCCESS);
        }
        return 1;  // Certificate chain is valid
    }
    // Verification failed - certificate chain is invalid
    return 0;
}

// Try to load system CA certificates from common directory locations
static int loadSystemCACertificates(HITLS_Config *config) {
    // Common CA certificate directory locations on different Linux distributions
    const char *caDirs[] = {
        "/etc/ssl/certs",                          // Debian/Ubuntu (contains individual cert files)
        "/etc/pki/tls/certs",                      // RHEL/CentOS/Fedora
        "/etc/pki/ca-trust/extracted/pem",         // CentOS 7+
        "/etc/ssl",                                // Alpine, macOS
        NULL
    };

    // Try to load from directory first
    for (int i = 0; caDirs[i] != NULL; i++) {
        int ret = HITLS_CFG_LoadVerifyDir(config, caDirs[i]);
        if (ret == HITLS_SUCCESS) {
            return HITLS_SUCCESS;
        }
    }

    // Fallback to bundle files if directory loading fails
    const char *caFiles[] = {
        "/etc/ssl/certs/ca-certificates.crt",      // Debian/Ubuntu bundle
        "/etc/pki/tls/certs/ca-bundle.crt",        // RHEL/CentOS/Fedora bundle
        "/etc/ssl/ca-bundle.pem",                  // OpenSUSE bundle
        "/etc/pki/ca-trust/extracted/pem/tls-ca-bundle.pem",  // CentOS 7+ bundle
        "/etc/ssl/cert.pem",                       // Alpine, macOS bundle
        NULL
    };

    for (int i = 0; caFiles[i] != NULL; i++) {
        int ret = HITLS_CFG_LoadVerifyFile(config, caFiles[i]);
        if (ret == HITLS_SUCCESS) {
            return HITLS_SUCCESS;
        }
    }

    return -1;
}

JNIEXPORT void JNICALL Java_org_openhitls_core_CryptoNative_sslCtxSetTrustAll
  (JNIEnv *env, jclass cls, jlong ctxPtr, jboolean trustAll) {
    HITLS_Config *config = (HITLS_Config *)ctxPtr;
    if (config == NULL) {
        return;
    }

    if (trustAll) {
        // Trust-all mode: Load CA certs for chain building but skip verification
        loadSystemCACertificates(config);

        // Enable verify-none mode (skip certificate verification)
        (void)HITLS_CFG_SetVerifyNoneSupport(config, true);

        // Disable client verify support for trust-all mode
        (void)HITLS_CFG_SetClientVerifySupport(config, false);

        // Set trust-all callback as a fallback
        HITLS_CFG_SetVerifyCb(config, trustAllVerifyCallback);
    } else {
        // Real verification mode: Load system CA certs and enable verification
        loadSystemCACertificates(config);

        // Disable verify-none mode (enable certificate verification)
        (void)HITLS_CFG_SetVerifyNoneSupport(config, false);

        // Set real verification callback
        HITLS_CFG_SetVerifyCb(config, realVerifyCallback);
    }
}

// Enable real certificate verification with system CA certificates
JNIEXPORT void JNICALL Java_org_openhitls_core_CryptoNative_sslCtxEnableVerification
  (JNIEnv *env, jclass cls, jlong ctxPtr) {
    HITLS_Config *config = (HITLS_Config *)ctxPtr;
    if (config == NULL) {
        return;
    }

    // Load system CA certificates
    loadSystemCACertificates(config);

    // Disable verify-none mode (enable certificate verification)
    (void)HITLS_CFG_SetVerifyNoneSupport(config, false);

    // Set real verification callback
    HITLS_CFG_SetVerifyCb(config, realVerifyCallback);
}

JNIEXPORT jint JNICALL Java_org_openhitls_core_CryptoNative_sslCtxLoadVerifyFile
  (JNIEnv *env, jclass cls, jlong ctxPtr, jstring filePath) {
    HITLS_Config *config = (HITLS_Config *)ctxPtr;
    if (config == NULL || filePath == NULL) {
        return -1;
    }

    const char *path = (*env)->GetStringUTFChars(env, filePath, NULL);
    if (path == NULL) {
        return -1;
    }

    int ret = HITLS_CFG_LoadVerifyFile(config, path);
    //fprintf(stderr, "sslCtxLoadVerifyFile: %s -> %d\n", path, ret);

    (*env)->ReleaseStringUTFChars(env, filePath, path);
    return ret;
}

// ==================== SSL Session Management ====================

JNIEXPORT jlong JNICALL Java_org_openhitls_core_CryptoNative_sslNew
  (JNIEnv *env, jclass cls, jlong ctxPtr) {
    HITLS_Config *config = (HITLS_Config *)ctxPtr;
    if (config == NULL) {
        throwSSLException(env, "SSL context is null");
        return 0;
    }

    SSLWrapper *wrapper = (SSLWrapper *)malloc(sizeof(SSLWrapper));
    if (wrapper == NULL) {
        throwSSLException(env, "Failed to allocate SSL wrapper");
        return 0;
    }
    memset(wrapper, 0, sizeof(SSLWrapper));

    wrapper->config = config;
    wrapper->ctx = HITLS_New(config);
    if (wrapper->ctx == NULL) {
        free(wrapper);
        throwSSLException(env, "Failed to create SSL session");
        return 0;
    }

    // Copy verify callback and verification settings from config to ctx if set
    HITLS_VerifyCb verifyCb = HITLS_CFG_GetVerifyCb(config);
    if (verifyCb != NULL) {
        (void)HITLS_SetVerifyCb(wrapper->ctx, verifyCb);

        // Also set verify-none and skip client verify on ctx level for trust-all mode
        (void)HITLS_SetVerifyNoneSupport(wrapper->ctx, true);
        (void)HITLS_SetClientVerifySupport(wrapper->ctx, false);
    }

    // Create custom Java BIO for TLS I/O (with separate input/output buffers)
    const BSL_UIO_Method *bioMethod = getJavaBioMethod();
    if (bioMethod == NULL) {
        HITLS_Free(wrapper->ctx);
        free(wrapper);
        throwSSLException(env, "Failed to create Java BIO method");
        return 0;
    }

    wrapper->uio = BSL_UIO_New(bioMethod);
    if (wrapper->uio == NULL) {
        HITLS_Free(wrapper->ctx);
        free(wrapper);
        throwSSLException(env, "Failed to create memory BIO");
        return 0;
    }

    // Set the UIO to the TLS context
    int ret = HITLS_SetUio(wrapper->ctx, wrapper->uio);
    if (ret != HITLS_SUCCESS) {
        BSL_UIO_Free(wrapper->uio);
        HITLS_Free(wrapper->ctx);
        free(wrapper);
        throwSSLException(env, "Failed to set UIO on TLS context");
        return 0;
    }

    wrapper->isClient = true;
    wrapper->handshakeState = 0;

    return (jlong)wrapper;
}

JNIEXPORT void JNICALL Java_org_openhitls_core_CryptoNative_sslFree
  (JNIEnv *env, jclass cls, jlong sslPtr) {
    if (sslPtr != 0) {
        SSLWrapper *wrapper = (SSLWrapper *)sslPtr;
        if (wrapper->ctx != NULL) {
            HITLS_Free(wrapper->ctx);
        }
        if (wrapper->uio != NULL) {
            BSL_UIO_Free(wrapper->uio);
        }
        free(wrapper);
    }
}

JNIEXPORT void JNICALL Java_org_openhitls_core_CryptoNative_sslSetConnectState
  (JNIEnv *env, jclass cls, jlong sslPtr, jint mode) {
    SSLWrapper *wrapper = (SSLWrapper *)sslPtr;
    if (wrapper != NULL) {
        wrapper->isClient = (mode == 1);  // 1 = client mode
        if (wrapper->ctx != NULL) {
            (void)HITLS_SetEndPoint(wrapper->ctx, wrapper->isClient);
        }
    }
}

JNIEXPORT void JNICALL Java_org_openhitls_core_CryptoNative_sslSetHostname
  (JNIEnv *env, jclass cls, jlong sslPtr, jstring hostname) {
    SSLWrapper *wrapper = (SSLWrapper *)sslPtr;
    if (wrapper == NULL || wrapper->ctx == NULL || hostname == NULL) {
        return;
    }

    const char *hostnameStr = (*env)->GetStringUTFChars(env, hostname, NULL);
    if (hostnameStr == NULL) {
        return;
    }

    // Set SNI hostname on the TLS context
    int ret = HITLS_SetServerName(wrapper->ctx, (uint8_t *)hostnameStr, strlen(hostnameStr));
    if (ret != HITLS_SUCCESS) {
        //fprintf(stderr, "HITLS_SetServerName failed: %d\n", ret);
    }

    (*env)->ReleaseStringUTFChars(env, hostname, hostnameStr);
}

// Last error for debugging
static int g_lastHandshakeError = 0;

// REC module return codes (need to be handled specially)
#define HITLS_REC_NORMAL_IO_BUSY_VAL           0x020A000A
#define HITLS_REC_NORMAL_RECV_BUF_EMPTY_VAL    0x020A000B
#define HITLS_REC_ERR_RECV_UNEXPECTED_MSG_VAL  0x020A0013

JNIEXPORT jint JNICALL Java_org_openhitls_core_CryptoNative_sslDoHandshake
  (JNIEnv *env, jclass cls, jlong sslPtr) {
    SSLWrapper *wrapper = (SSLWrapper *)sslPtr;
    if (wrapper == NULL || wrapper->ctx == NULL) {
        return -1;  // TLS_ERROR
    }

    // Set current ctx for verify callback (used if trust-all callback is set)
    g_currentCtxForVerify = wrapper->ctx;

    // Use HITLS_Connect for client, HITLS_Accept for server
    int ret;
    if (wrapper->isClient) {
        ret = HITLS_Connect(wrapper->ctx);
    } else {
        ret = HITLS_Accept(wrapper->ctx);
    }

    // Clear current ctx after handshake step
    g_currentCtxForVerify = NULL;

    // Store for debugging
    g_lastHandshakeError = ret;

    // Handle HITLS return codes (defined in hitls_error.h)
    // Simple return codes: HITLS_SUCCESS = 0, HITLS_WANT_CONNECT = 1, HITLS_WANT_ACCEPT = 2,
    // HITLS_WANT_READ = 3, HITLS_WANT_WRITE = 4
    // REC module codes: HITLS_REC_NORMAL_IO_BUSY = 0x020A000A, HITLS_REC_NORMAL_RECV_BUF_EMPTY = 0x020A000B

    if (ret == HITLS_SUCCESS) {
        wrapper->handshakeState = 2;  // Complete
        return 0;  // TLS_SUCCESS
    }

    // Check for "want" conditions
    if (ret == HITLS_WANT_CONNECT || ret == HITLS_WANT_ACCEPT || ret == HITLS_WANT_WRITE) {
        wrapper->handshakeState = 1;
        return 2;  // TLS_WANT_WRITE
    }

    if (ret == HITLS_WANT_READ) {
        wrapper->handshakeState = 1;
        return 1;  // TLS_WANT_READ
    }

    // Check for REC module "normal" conditions (need more I/O)
    if (ret == (int)HITLS_REC_NORMAL_RECV_BUF_EMPTY_VAL) {
        wrapper->handshakeState = 1;
        return 1;  // TLS_WANT_READ - need more data from network
    }

    if (ret == (int)HITLS_REC_NORMAL_IO_BUSY_VAL) {
        wrapper->handshakeState = 1;
        return 2;  // TLS_WANT_WRITE - need to flush output
    }

    // REC_ERR_RECV_UNEXPECTED_MSG can occur when unwrap is called with no data
    // (e.g., server waiting for ClientHello but nothing received yet)
    if (ret == (int)HITLS_REC_ERR_RECV_UNEXPECTED_MSG_VAL) {
        wrapper->handshakeState = 1;
        return 1;  // TLS_WANT_READ - need more data from network
    }

    // Error occurred
    //fprintf(stderr, "HITLS handshake error: %d (0x%x)\n", ret, ret);
    return -1;  // TLS_ERROR
}

// Get the last handshake error code for debugging
JNIEXPORT jint JNICALL Java_org_openhitls_core_CryptoNative_sslGetLastError
  (JNIEnv *env, jclass cls, jlong sslPtr) {
    return g_lastHandshakeError;
}

JNIEXPORT jint JNICALL Java_org_openhitls_core_CryptoNative_sslRead
  (JNIEnv *env, jclass cls, jlong sslPtr, jbyteArray buffer, jint offset, jint length) {
    SSLWrapper *wrapper = (SSLWrapper *)sslPtr;
    if (wrapper == NULL || wrapper->ctx == NULL) {
        return -1;
    }

    jbyte *buf = (*env)->GetByteArrayElements(env, buffer, NULL);
    if (buf == NULL) {
        return -1;
    }

    uint32_t readLen = 0;
    int ret = HITLS_Read(wrapper->ctx, (uint8_t *)(buf + offset), length, &readLen);

    (*env)->ReleaseByteArrayElements(env, buffer, buf, 0);

    if (ret == HITLS_SUCCESS) {
        return (jint)readLen;
    } else if (ret == HITLS_WANT_READ) {
        return 1;  // TLS_WANT_READ
    } else if (ret == HITLS_WANT_WRITE) {
        return 2;  // TLS_WANT_WRITE
    }

    return -1;
}

JNIEXPORT jint JNICALL Java_org_openhitls_core_CryptoNative_sslWrite
  (JNIEnv *env, jclass cls, jlong sslPtr, jbyteArray data, jint offset, jint length) {
    SSLWrapper *wrapper = (SSLWrapper *)sslPtr;
    if (wrapper == NULL || wrapper->ctx == NULL) {
        return -1;
    }

    jbyte *buf = (*env)->GetByteArrayElements(env, data, NULL);
    if (buf == NULL) {
        return -1;
    }

    uint32_t writeLen = 0;
    int ret = HITLS_Write(wrapper->ctx, (uint8_t *)(buf + offset), length, &writeLen);

    (*env)->ReleaseByteArrayElements(env, data, buf, JNI_ABORT);

    if (ret == HITLS_SUCCESS) {
        return (jint)writeLen;
    } else if (ret == HITLS_WANT_WRITE) {
        return 2;  // TLS_WANT_WRITE
    } else if (ret == HITLS_WANT_READ) {
        return 1;  // TLS_WANT_READ
    }

    return -1;
}

JNIEXPORT jint JNICALL Java_org_openhitls_core_CryptoNative_sslShutdown
  (JNIEnv *env, jclass cls, jlong sslPtr) {
    SSLWrapper *wrapper = (SSLWrapper *)sslPtr;
    if (wrapper != NULL && wrapper->ctx != NULL) {
        return HITLS_Close(wrapper->ctx);
    }
    return 0;
}

// ==================== Memory BIO Operations ====================

// Write network data TO TLS (incoming encrypted data from network)
JNIEXPORT jint JNICALL Java_org_openhitls_core_CryptoNative_sslBioWriteNetData
  (JNIEnv *env, jclass cls, jlong sslPtr, jbyteArray data, jint offset, jint length) {
    SSLWrapper *wrapper = (SSLWrapper *)sslPtr;
    if (wrapper == NULL || wrapper->uio == NULL) {
        return -1;
    }

    JavaBioData *bioData = (JavaBioData *)BSL_UIO_GetCtx(wrapper->uio);
    if (bioData == NULL) {
        return -1;
    }

    jbyte *buf = (*env)->GetByteArrayElements(env, data, NULL);
    if (buf == NULL) {
        return -1;
    }

    // Compact input buffer if needed
    if (bioData->inBufPos > 0 && bioData->inBufPos >= bioData->inBufLen) {
        bioData->inBufPos = 0;
        bioData->inBufLen = 0;
    }

    // Check if we need to expand the input buffer
    size_t needed = bioData->inBufLen + (size_t)length;
    if (needed > bioData->inBufSize) {
        size_t newSize = bioData->inBufSize * 2;
        if (newSize < needed) newSize = needed;
        if (newSize < 32768) newSize = 32768;
        uint8_t *newBuf = (uint8_t *)realloc(bioData->inBuf, newSize);
        if (newBuf == NULL) {
            (*env)->ReleaseByteArrayElements(env, data, buf, JNI_ABORT);
            return -1;
        }
        bioData->inBuf = newBuf;
        bioData->inBufSize = newSize;
    }

    // Copy data to input buffer
    memcpy(bioData->inBuf + bioData->inBufLen, buf + offset, (size_t)length);
    bioData->inBufLen += (size_t)length;

    (*env)->ReleaseByteArrayElements(env, data, buf, JNI_ABORT);
    return length;
}

// Read network data FROM TLS (outgoing encrypted data to send over network)
JNIEXPORT jint JNICALL Java_org_openhitls_core_CryptoNative_sslBioReadNetData
  (JNIEnv *env, jclass cls, jlong sslPtr, jbyteArray buffer, jint offset, jint length) {
    SSLWrapper *wrapper = (SSLWrapper *)sslPtr;
    if (wrapper == NULL || wrapper->uio == NULL) {
        return -1;
    }

    JavaBioData *bioData = (JavaBioData *)BSL_UIO_GetCtx(wrapper->uio);
    if (bioData == NULL) {
        return -1;
    }

    size_t available = bioData->outBufLen - bioData->outBufPos;
    if (available == 0) {
        return 0;  // No data available
    }

    jbyte *buf = (*env)->GetByteArrayElements(env, buffer, NULL);
    if (buf == NULL) {
        return -1;
    }

    size_t toRead = (size_t)length < available ? (size_t)length : available;
    memcpy(buf + offset, bioData->outBuf + bioData->outBufPos, toRead);
    bioData->outBufPos += toRead;

    // Compact output buffer if fully consumed
    if (bioData->outBufPos >= bioData->outBufLen) {
        bioData->outBufPos = 0;
        bioData->outBufLen = 0;
    }

    (*env)->ReleaseByteArrayElements(env, buffer, buf, 0);
    return (jint)toRead;
}

// Get pending write data (outgoing network data waiting to be read by Java)
JNIEXPORT jint JNICALL Java_org_openhitls_core_CryptoNative_sslBioPendingWrite
  (JNIEnv *env, jclass cls, jlong sslPtr) {
    SSLWrapper *wrapper = (SSLWrapper *)sslPtr;
    if (wrapper == NULL || wrapper->uio == NULL) {
        //fprintf(stderr, "sslBioPendingWrite: wrapper or uio is NULL\n");
        return 0;
    }

    JavaBioData *bioData = (JavaBioData *)BSL_UIO_GetCtx(wrapper->uio);
    if (bioData == NULL) {
        //fprintf(stderr, "sslBioPendingWrite: bioData is NULL\n");
        return 0;
    }

    jint pending = (jint)(bioData->outBufLen - bioData->outBufPos);
    return pending;
}

// Get pending read data (incoming network data waiting for HITLS to process)
JNIEXPORT jint JNICALL Java_org_openhitls_core_CryptoNative_sslBioPendingRead
  (JNIEnv *env, jclass cls, jlong sslPtr) {
    SSLWrapper *wrapper = (SSLWrapper *)sslPtr;
    if (wrapper == NULL || wrapper->uio == NULL) {
        return 0;
    }

    JavaBioData *bioData = (JavaBioData *)BSL_UIO_GetCtx(wrapper->uio);
    if (bioData == NULL) {
        return 0;
    }

    return (jint)(bioData->inBufLen - bioData->inBufPos);
}

// ==================== Session Information ====================

JNIEXPORT jint JNICALL Java_org_openhitls_core_CryptoNative_sslGetVersion
  (JNIEnv *env, jclass cls, jlong sslPtr) {
    SSLWrapper *wrapper = (SSLWrapper *)sslPtr;
    if (wrapper == NULL || wrapper->ctx == NULL) {
        return 0;
    }

    uint16_t version = 0;
    HITLS_GetNegotiatedVersion(wrapper->ctx, &version);
    return (jint)version;
}

JNIEXPORT jstring JNICALL Java_org_openhitls_core_CryptoNative_sslGetCipherSuite
  (JNIEnv *env, jclass cls, jlong sslPtr) {
    return (*env)->NewStringUTF(env, "TLS_AES_128_GCM_SHA256");  // Default
}

JNIEXPORT jobjectArray JNICALL Java_org_openhitls_core_CryptoNative_sslGetPeerCertificateChain
  (JNIEnv *env, jclass cls, jlong sslPtr) {
    return NULL;  // No peer certificates
}

JNIEXPORT jbyteArray JNICALL Java_org_openhitls_core_CryptoNative_sslGetSessionId
  (JNIEnv *env, jclass cls, jlong sslPtr) {
    jbyteArray sessionId = (*env)->NewByteArray(env, 32);
    return sessionId;  // Empty session ID
}

JNIEXPORT jstring JNICALL Java_org_openhitls_core_CryptoNative_sslGetErrorString
  (JNIEnv *env, jclass cls, jlong sslPtr) {
    return (*env)->NewStringUTF(env, "");
}

JNIEXPORT jint JNICALL Java_org_openhitls_core_CryptoNative_sslGetError
  (JNIEnv *env, jclass cls, jlong sslPtr) {
    return 0;  // No error
}

// Get negotiated key exchange group
JNIEXPORT jint JNICALL Java_org_openhitls_core_CryptoNative_sslGetNegotiatedGroup
  (JNIEnv *env, jclass cls, jlong sslPtr) {
    SSLWrapper *wrapper = (SSLWrapper *)sslPtr;
    if (wrapper == NULL || wrapper->ctx == NULL) {
        return 0;
    }

    uint16_t group = 0;
    int ret = HITLS_GetNegotiateGroup(wrapper->ctx, &group);
    if (ret != HITLS_SUCCESS) {
        return 0;
    }
    return (jint)group;
}

// Set supported key exchange groups
JNIEXPORT void JNICALL Java_org_openhitls_core_CryptoNative_sslCtxSetGroups
  (JNIEnv *env, jclass cls, jlong ctxPtr, jintArray groups) {
    HITLS_Config *config = (HITLS_Config *)ctxPtr;
    if (config == NULL || groups == NULL) {
        return;
    }

    jsize len = (*env)->GetArrayLength(env, groups);
    if (len <= 0) {
        return;
    }

    jint *groupArray = (*env)->GetIntArrayElements(env, groups, NULL);
    if (groupArray == NULL) {
        return;
    }

    // Convert jint array to uint16_t array
    uint16_t *groupList = (uint16_t *)malloc(len * sizeof(uint16_t));
    if (groupList == NULL) {
        (*env)->ReleaseIntArrayElements(env, groups, groupArray, 0);
        return;
    }

    for (jsize i = 0; i < len; i++) {
        groupList[i] = (uint16_t)groupArray[i];
    }

    HITLS_CFG_SetGroups(config, groupList, (uint32_t)len);

    free(groupList);
    (*env)->ReleaseIntArrayElements(env, groups, groupArray, 0);
}
